package benchmark.energy.client;

import benchmark.client.DBProxy;
import benchmark.energy.EnergyDataTxGenerator;
import benchmark.energy.tx.*;
import benchmark.transaction.definition.AbstractTransaction;
import benchmark.utils.*;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;

import java.io.IOException;
import java.sql.*;
import java.util.*;
import java.util.concurrent.*;

public class PostgreSQLExecutorClient implements DBProxy {
    private ThreadPoolExecutor exe;
    private BlockingQueue<Connection> connectionPool = new LinkedBlockingQueue<>();
    private ListeningExecutorService service;
    private Map<Connection, Integer> connIdMap = new HashMap<>();
    private HashMap<Integer, String> id2Name = new HashMap<>();
    private HashMap<String, Integer> name2Id = new HashMap<>();

    public PostgreSQLExecutorClient(String serverHost, int parallelCnt, int queueLength) throws IOException, ClassNotFoundException, SQLException, InterruptedException {
        Class.forName("org.postgresql.Driver");
        String initURL = "jdbc:postgresql://" + serverHost + ":5432/";
        String dbURL = initURL + "beijing_traffic?reWriteBatchedInserts=true";
        Connection initConn = DriverManager.getConnection(initURL, "postgres", "root");
        createDB(initConn);
        initConn.close();
        this.exe = new ThreadPoolExecutor(parallelCnt, parallelCnt, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(queueLength), (r, executor) -> {
            if (!executor.isShutdown()) {
                try {
                    executor.getQueue().put(r);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        });
        exe.prestartAllCoreThreads();
        this.service = MoreExecutors.listeningDecorator(exe);
        for (int i = 0; i < parallelCnt; i++) {
            Connection conn = DriverManager.getConnection(dbURL, "postgres", "root");
            this.connectionPool.offer(conn);
            connIdMap.put(conn, i);
        }
    }

    private void createDB(Connection conn) {
        try (Statement stmt = conn.createStatement()) {
            conn.setAutoCommit(true);
            ResultSet rs = stmt.executeQuery("select database.datname from pg_catalog.pg_database database where database.datname = 'beijing_traffic'");
            if (rs.next()) return;
            stmt.execute("create database beijing_traffic");
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    private void createTables() throws InterruptedException {
        Connection conn = connectionPool.take();
        try (Statement stmt = conn.createStatement()) {
            conn.setAutoCommit(true);
            stmt.execute("CREATE table energy_node(raw_id int PRIMARY KEY, name varchar(255), country varchar(255), volt int, lat float, lon float");
            stmt.execute("CREATE table electric_line(r_start int, r_end int, x float, y float, r_limit int, r_length float)");
            stmt.execute("CREATE table temporal_status(node_id int, t int, load_signal FLOAT, " +
                            "solar_cosmo_p FLOAT, solar_cosmo_u FLOAT, solar_ecmwf_p FLOAT, solar_ecmwf_u FLOAT, " +
                            "wind_cosmo_p FLOAT, wind_cosmo_u FLOAT, wind_ecmwf_p FLOAT, wind_ecmwf_u FLOAT )");
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            connectionPool.put(conn);
        }
    }

    public void createIndexes() throws InterruptedException {
        Connection conn = connectionPool.take();
        try (Statement stmt = conn.createStatement()) {
            conn.setAutoCommit(true);
            stmt.execute("create index ind_id_t on temporal_status(node_id, t)");
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            connectionPool.put(conn);
        }
    }

    @Override
    public ListenableFuture<ServerResponse> execute(AbstractTransaction tx) {
        switch (tx.getTxType()) {
            case TX_IMPORT_STATIC_DATA:
                return this.service.submit(execute((ImportStaticDataTx) tx));
            case TX_IMPORT_TEMPORAL_DATA:
                return this.service.submit(execute((ImportTemporalDataTx) tx));
            case TX_QUERY_SNAPSHOT:
                return this.service.submit(execute((SnapshotQueryTx) tx));
            case TX_QUERY_SNAPSHOT_AGGR_MAX:
                return this.service.submit(execute((SnapshotAggrMaxTx) tx));
            case TX_QUERY_SNAPSHOT_AGGR_DURATION:
                return this.service.submit(execute((SnapshotAggrDurationTx) tx));
            case TX_ENTITY_TEMPORAL_CONDITION_QUERY:
                return this.service.submit(execute((EntityTemporalConditionTx) tx));
//            case TX_UPDATE_TEMPORAL_DATA:
//                return this.service.submit(execute((UpdateTemporalDataTx) tx));
            default:
                throw new UnsupportedOperationException();
        }
    }

    @Override
    public void createDB() throws IOException {

    }

    @Override
    public void restartDB() throws IOException {

    }

    @Override
    public void shutdownDB() throws IOException {

    }

    @Override
    public void close() throws IOException, InterruptedException {
        try {
            service.shutdown();
            while (!service.isTerminated()) {
                service.awaitTermination(10, TimeUnit.SECONDS);
                long completeCnt = exe.getCompletedTaskCount();
                int remains = exe.getQueue().size();
                System.out.println(completeCnt + " / " + (completeCnt + remains) + " query completed.");
            }
            while (!connectionPool.isEmpty()) {
                Connection conn = connectionPool.take();
                conn.close();
            }
            System.out.println("Client exit. send " + exe.getCompletedTaskCount() + " lines.");
        } catch (SQLException e) {
            e.printStackTrace();
            throw new IOException(e);
        }
    }

    @Override
    public String testServerClientCompatibility() {
        return null;
    }

//    private boolean buildMapBetweenIdAndName() throws InterruptedException {
//        boolean ret = false;
//        Connection conn = connectionPool.take();
//        ret = buildMapBetweenIdAndName(conn);
//        connectionPool.put(conn);
//        return ret;
//    }
//
//    private boolean buildMapBetweenIdAndName(Connection conn) {
//        try (PreparedStatement stmt0 = conn.prepareStatement("select r_id, r_address from road");
//             PreparedStatement stmt1 = conn.prepareStatement("select relname from pg_class where relname = 'road'")) {
//            ResultSet r = stmt1.executeQuery();
//            if (!r.next()) return false;
//            ResultSet rs = stmt0.executeQuery();
//            while (rs.next()) {
//                Long id = rs.getLong("r_id");
//                String name = rs.getString("r_address");
//                id2Name.put(id, name);
//                name2Id.put(name, id);
//            }
//        } catch (SQLException e) {
//            e.printStackTrace();
//        }
//        return true;
//    }

    private Callable<ServerResponse> execute(ImportStaticDataTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                try (PreparedStatement stat1 = conn.prepareStatement("INSERT INTO energy_node VALUES (?,?,?,?,?)");
                     PreparedStatement stat2 = conn.prepareStatement("INSERT INTO electric_line VALUES (?, ?, ?, ?, ?, ?)")) {
                    conn.setAutoCommit(false);
                    for (ImportStaticDataTx.EnergyNode p : tx.getNodes()) {
                        stat1.setInt(1, p.getRawId());
                        stat1.setString(2, p.getName());
                        stat1.setString(3, p.getCountry());
                        stat1.setInt(4, p.getVoltage());
                        stat1.setFloat(5, p.getLatitude());
                        stat1.setFloat(6, p.getLongitude());
                        stat1.addBatch();
                    }
                    stat1.executeBatch();
                    for (ImportStaticDataTx.EnergyRel r : tx.getRels()) {
                        stat2.setInt(1, r.getFromRawId());
                        stat2.setInt(2, r.getToRawId());
                        stat2.setFloat(3, r.getX());
                        stat2.setFloat(4, r.getY());
                        stat2.setInt(6, r.getLimit());
                        stat2.setFloat(7, r.getLength());
                        stat2.addBatch();
                    }
                    stat2.executeBatch();
                    conn.commit();
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                return new AbstractTransaction.Result();
            }
        };
    }

    private Callable<ServerResponse> execute(ImportTemporalDataTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                try (PreparedStatement stmt = conn.prepareStatement("insert into temporal_status values(?,?,?,?,?,?,?,?,?)")) {
                    conn.setAutoCommit(false);
                    for (ImportTemporalDataTx.EnergyStatus s : tx.data) {
                        stmt.setLong(1, s.getNodeId());
                        stmt.setInt(2, s.getTimePoint());
                        List<Float> ss = s.getStatus();
                        for(int i=3; i<3+ss.size();i++){
                            stmt.setFloat(i, ss.get(i-3));
                        }
                        stmt.addBatch();
                    }
                    stmt.executeBatch();
                    conn.commit();
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                return new AbstractTransaction.Result();
            }
        };
    }

    private Callable<ServerResponse> execute(UpdateTemporalDataTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                conn.setAutoCommit(false);
                int nodeId = name2Id.get(tx.getEntityId());
                String[] props = EnergyDataTxGenerator.properties;
                String propStr = String.join(",", props);
                try (PreparedStatement stmt1 = conn.prepareStatement("select node_id, max(t) as max_t from temporal_status where node_id=? and t<=? group by node_id");
                     PreparedStatement stmt2 = conn.prepareStatement("select "+propStr+" from temporal_status where node_id=? and t=?");
                     PreparedStatement stmt3 = conn.prepareStatement("insert into temporal_status(t, node_id, "+propStr+") values(?,?,?,?,?,?,?,?,?,?,?)");
                     PreparedStatement stmt4 = conn.prepareStatement("delete from temporal_status where node_id=? and t>=? and t<=?")) {
                    stmt1.setInt(1, nodeId);
                    stmt1.setInt(2, tx.getEndTime());
                    ResultSet rs = stmt1.executeQuery();

                    // insert into table from t1 + 1 with the old value.
                    if (rs.next()) {
                        stmt2.setInt(1, rs.getInt("node_id"));
                        stmt2.setInt(2, rs.getInt("max_t"));
                        ResultSet r = stmt2.executeQuery();
                        if (r.next()) {
                            stmt3.setInt(1, tx.getEndTime() + 1);
                            stmt3.setInt(2, nodeId);
                            for(int i=3; i<props.length; i++){
                                stmt3.setFloat(i, r.getFloat(props[i-3]));
                            }
                            stmt3.execute();
                        }
                    }
                    // delete the [t0, t1] rows.
                    stmt4.setInt(1, nodeId);
                    stmt4.setInt(2, tx.getStartTime());
                    stmt4.setInt(3, tx.getEndTime());
                    stmt4.execute();
                    // insert into table from t0 with the new value.
                    stmt3.setInt(1, tx.getStartTime());
                    stmt3.setInt(2, nodeId);
                    for(int i=3; i<props.length; i++){
                        stmt3.setFloat(i, tx.getValues().get(i-3));
                    }
                    stmt3.execute();
                    conn.commit();
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                return new AbstractTransaction.Result();
            }
        };
    }

    private Callable<ServerResponse> execute(SnapshotQueryTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<Pair<String, Float>> res = new ArrayList<>();
                try (PreparedStatement stmt1 = conn.prepareStatement("select name, " + tx.getPropertyName() + " as property from temporal_status, " +
                             "(select node_id, max(t) as max_t from temporal_status where t <= ? group by node_id) as tmp where temporal_status.node_id = tmp.node_id and temporal_status.t = tmp.max_t")) {
                    conn.setAutoCommit(false);
                    int snapshotTime = tx.getTimestamp();
                    stmt1.setInt(1, snapshotTime);
                    ResultSet rs = stmt1.executeQuery();
                    while (rs.next()) {
                        res.add(Pair.of(String.valueOf(rs.getInt("node_id")), rs.getFloat("property")));
                    }
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
                result.setStatus(res);
                return result;
            }
        };
    }

    private Callable<ServerResponse> execute(SnapshotAggrMaxTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<Pair<String, Float>> res = new ArrayList<>();
                try (PreparedStatement stmt2 = conn.prepareStatement("select node_id, max(" + tx.getP() + ") as max_p from temporal_status, " +
                        "(select node_id, max(t) as max_t from temporal_status where t<=? group by node_id) as tmp " +
                        "where tmp.node_id=temporal_status.node_id and tmp.max_t<=temporal_status.t and temporal_status.t<=? group by node_id")) {
                    conn.setAutoCommit(false);
                    stmt2.setInt(1, tx.getT0());
                    stmt2.setInt(2, tx.getT1());
                    ResultSet rs = stmt2.executeQuery();
                    HashMap<String, Integer> st = new HashMap<>();
                    while (rs.next()) {
                        res.add(Pair.of(String.valueOf(rs.getInt("node_id")), rs.getFloat("max_p")));
                    }
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
                result.setPropMaxValue(res);
                return result;
            }
        };
    }

    private Callable<ServerResponse> execute(SnapshotAggrDurationTx tx) {
        final TreeSet<Integer> intStart = tx.getIntervalStarts();
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<Triple<String, Integer, Integer>> res = new ArrayList<>();
                try (PreparedStatement stmt0 = conn.prepareStatement("select node_id, max(t) as real_st from temporal_status where t <= ? group by node_id");
                     PreparedStatement stmt1 = conn.prepareStatement("select t, " + tx.getP() + " as property from temporal_status where node_id = ? and t >= ? and t <= ? order by t")) {
                    conn.setAutoCommit(false);
                    stmt0.setInt(1, tx.getT0());
                    ResultSet rs = stmt0.executeQuery();
                    while (rs.next()) {
                        int id = rs.getInt("node_id");
                        stmt1.setLong(1, id);
                        stmt1.setInt(2, rs.getInt("real_st"));
                        stmt1.setInt(3, tx.getT1());
                        ResultSet r = stmt1.executeQuery();
                        int lastTime = -1;
                        float lastStatus = Float.MIN_VALUE;
                        HashMap<Integer, Integer> buffer = new HashMap<>();
                        while (r.next()) {
                            float property = r.getFloat("property");
                            int t = r.getInt("t");
                            if (lastTime != -1) {
                                int duration = t - lastTime;
                                Integer statusGrp = intStart.floor((int) lastStatus);
                                if(statusGrp!=null) buffer.merge(statusGrp, duration, Integer::sum);
                                lastTime = t;
                            } else {
                                lastTime = tx.getT0();
                            }
                            lastStatus = property;
                        }
                        if (lastTime != -1) {
                            Integer statusGrp = intStart.floor((int) lastStatus);
                            if(statusGrp!=null) buffer.merge(statusGrp, tx.getT1() - lastTime, Integer::sum);
                        }
                        for (Map.Entry<Integer, Integer> e : buffer.entrySet()) {
                            res.add(Triple.of(id2Name.get(id), e.getKey(), e.getValue()));
                        }
                    }
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
                result.setStatusDuration(res);
                return result;
            }
        };
    }

    private Callable<ServerResponse> execute(EntityTemporalConditionTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<String> res = new ArrayList<>();
                try (PreparedStatement stmt0 = conn.prepareStatement("select node_id from temporal_status as ts, " +
                        "(select node_id, max(t) as max_t from temporal_status where t<=? group by node_id) as tmp " +
                        "WHERE tmp.node_id=ts.node_id AND ts."  + tx.getP() + " >= ? and ts." + tx.getP() + " <= ?"+
                        "tmp.max_t<=ts.t AND ts.t<=? LIMIT 1")) {
                    conn.setAutoCommit(false);
                    ResultSet rs = stmt0.executeQuery();
                    while (rs.next()) {
                        res.add(id2Name.get(rs.getInt("node_id")));
                    }
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
                result.setEntities(res);
                return result;
            }
        };
    }

    private abstract class Req implements Callable<ServerResponse> {
        private final TimeMonitor timeMonitor = new TimeMonitor();
        final AbstractTransaction.Metrics metrics = new AbstractTransaction.Metrics();

        private Req() {
            timeMonitor.begin("Wait in queue");
        }

        @Override
        public ServerResponse call() throws Exception {
            try {
                Connection conn = connectionPool.take();
                timeMonitor.mark("Wait in queue", "query");
                AbstractTransaction.Result result = executeQuery(conn);
                timeMonitor.end("query");
                if (result == null) throw new RuntimeException("[Got null. Server close connection]");
                connectionPool.put(conn);
                metrics.setWaitTime(Math.toIntExact(timeMonitor.duration("Wait in queue")));
                metrics.setSendTime(timeMonitor.beginT("query"));
                metrics.setExeTime(Math.toIntExact(timeMonitor.duration("query")));
                metrics.setConnId(connIdMap.get(conn));
                ServerResponse response = new ServerResponse();
                response.setMetrics(metrics);
                response.setResult(result);
                return response;
            } catch (Exception e) {
                e.printStackTrace();
                throw e;
            }
        }

        protected abstract AbstractTransaction.Result executeQuery(Connection conn) throws Exception;
    }
}
